/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [https://neo4j.com]
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.neo4j.cypher.internal.parser

import com.squareup.javapoet.ClassName
import com.squareup.javapoet.JavaFile
import com.squareup.javapoet.MethodSpec
import com.squareup.javapoet.TypeSpec
import org.antlr.v4.runtime.Parser
import org.antlr.v4.runtime.ParserRuleContext
import org.antlr.v4.runtime.tree.ParseTreeListener
import org.apache.maven.plugin.AbstractMojo
import org.apache.maven.plugins.annotations.LifecyclePhase
import org.apache.maven.plugins.annotations.Mojo
import org.apache.maven.plugins.annotations.Parameter
import org.apache.maven.plugins.annotations.ResolutionScope
import org.apache.maven.project.MavenProject
import org.neo4j.cypher.internal.parser.GenerateListenerMavenPlugin.ParserMeta
import org.neo4j.cypher.internal.parser.GenerateListenerMavenPlugin.RuleMeta

import java.io.File
import java.net.URLClassLoader

import scala.jdk.CollectionConverters.SeqHasAsJava

import javax.lang.model.element.Modifier

/**
 * Generates a custom optimised parser listener based on an antlr parser.
 */
@Mojo(
  name = "generate-antlr-listener",
  defaultPhase = LifecyclePhase.GENERATE_SOURCES,
  requiresDependencyResolution = ResolutionScope.COMPILE_PLUS_RUNTIME,
  requiresProject = false,
  threadSafe = true
)
class GenerateListenerMavenPlugin extends AbstractMojo {

  @Parameter(defaultValue = "${project}", required = true, readonly = true)
  protected var project: MavenProject = _

  @Parameter(defaultValue = "${project.build.directory}/generated-sources/parser")
  protected var outputDirectory: File = _

  @Parameter(required = true)
  protected var parserClass: String = _

  @Parameter(required = true)
  protected var listenerName: String = _

  override def execute(): Unit = {
    val parser = GenerateListenerMavenPlugin.parserMeta(parserClass, getClassLoader)
    val parserName = parser.cls.getSimpleName

    val parserListenerSpec = TypeSpec.interfaceBuilder(s"${parserName}Listener")
      .addModifiers(Modifier.PUBLIC)
      .addSuperinterface(classOf[ParseTreeListener])
      .addMethods(genExitMethods(parser).asJava)
      .addJavadoc(s"A ParseTreeListener for $parserName.\nGenerated by ${getClass.getCanonicalName}")
      .build
    JavaFile.builder(parser.cls.getPackageName, parserListenerSpec).build
      .writeTo(outputDirectory)

    val abstractParseListenerBuilderSpec = TypeSpec.classBuilder(listenerName)
      .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
      .addSuperinterface(ClassName.get(parser.cls.getPackageName, parserListenerSpec.name))
      .addMethod(genExitAllMethod(parser))
      .addJavadoc(s"Optimised implementation of ${parserListenerSpec.name}.\nGenerated by ${getClass.getCanonicalName}")
      .build
    JavaFile.builder(parser.cls.getPackageName, abstractParseListenerBuilderSpec).build
      .writeTo(outputDirectory)
  }

  /*
   * Generate exit methods for all rule context classes,
   * for example: void exitStatement(Parser.StatementContext ctx);
   */
  private def genExitMethods(meta: ParserMeta) = {
    meta.rules.map { rule =>
      MethodSpec.methodBuilder(rule.exitName)
        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
        .returns(classOf[Unit])
        .addParameter(rule.cls, "ctx")
        .build()
    }
  }

  /*
   * Generate the optimised exitEveryRule method.
   */
  private def genExitAllMethod(parser: ParserMeta) = {
    MethodSpec.methodBuilder("exitEveryRule")
      .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
      .addAnnotation(classOf[Override])
      .returns(classOf[Unit])
      .addParameter(classOf[ParserRuleContext], "ctx")
      .addJavadoc(
        """
          |Optimised exit method.
          |
          |This compiles into a blazingly fast tableswitch (jump table)
          |and has been shown to be faster than the generated listeners that antlr provides.
          |""".stripMargin
      )
      .addCode(
        s"""
           |switch (ctx.getRuleIndex()) {
           |${parser.rules.map(r => genCase(parser, r)).mkString("\n")}
           |  default -> throw new IllegalStateException("Unknown rule index " + ctx.getRuleIndex());
           |}
           |""".stripMargin.trim
      )
      .build()
  }

  private def genCase(parser: ParserMeta, rule: RuleMeta) = {
    val ruleIndexField = parser.cls.getSimpleName + "." + rule.indexFieldName // YourParser.RULE_theRule
    val ctxClassName = s"${parser.cls.getSimpleName}.${rule.cls.getSimpleName}" // YourParser.TheRuleContext
    s"  case $ruleIndexField -> ${rule.exitName}(($ctxClassName) ctx);"
  }

  private def getClassLoader: ClassLoader = {
    Option(project.getRuntimeClasspathElements) match {
      case Some(classpathElements) =>
        val urls = classpathElements.toArray(Array[String]()).map(e => new File(e).toURI.toURL)
        new URLClassLoader(urls, getClass.getClassLoader)
      case _ =>
        Thread.currentThread.getContextClassLoader
    }
  }
}

object GenerateListenerMavenPlugin {
  case class ParserMeta(cls: Class[_ <: Parser], rules: Seq[RuleMeta])

  case class RuleMeta(name: String, index: Int, cls: Class[_]) {
    def indexFieldName: String = s"RULE_$name"
    def exitName: String = s"exit${name.capitalize}"
  }

  private def parserMeta(parserClassName: String, classLoader: ClassLoader): ParserMeta = {
    val parserClass = classLoader.loadClass(parserClassName).asSubclass(classOf[Parser])
    val names: Array[String] = parserClass.getField("ruleNames").get(null).asInstanceOf[Array[String]]
    ParserMeta(
      cls = parserClass,
      rules = names.zipWithIndex.map { case (name, index) =>
        val className = parserClass.getName + "$" + name.capitalize + "Context"
        val ruleClass = parserClass.getClassLoader.loadClass(className)
        RuleMeta(name, index, ruleClass)
      }
    )
  }
}
